import torch
import torch.nn as nn
import torch.nn.functional as F
from .batchnorm import SynchronizedBatchNorm2d

BatchNorm2d = SynchronizedBatchNorm2d

class BasicConv2d(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_planes, out_planes,
                              kernel_size=kernel_size, stride=stride,
                              padding=padding, dilation=dilation, bias=False)
        self.bn = nn.BatchNorm2d(out_planes)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return x

class SpeckleNoiseReduction(nn.Module):
    """超声图像斑点噪声抑制模块"""
    def __init__(self, in_channels, out_channels):
        super(SpeckleNoiseReduction, self).__init__()
        self.conv1 = BasicConv2d(in_channels, out_channels, 3, padding=1)
        self.conv2 = BasicConv2d(out_channels, out_channels, 3, padding=1)
        self.conv3 = BasicConv2d(out_channels, out_channels, 3, padding=1)
        # 如果输入输出通道数不同，需要投影层
        if in_channels != out_channels:
            self.projection = BasicConv2d(in_channels, out_channels, 1)
        else:
            self.projection = None
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        # 多尺度去噪
        x1 = self.relu(self.conv1(x))
        x2 = self.relu(self.conv2(x1))
        x3 = self.relu(self.conv3(x2))
        
        # 残差连接：确保通道数匹配
        if self.projection is not None:
            x_res = self.projection(x)
        else:
            x_res = x
            
        return x_res + x3

class EdgeEnhancement(nn.Module):
    """边缘增强模块"""
    def __init__(self, in_channels, out_channels):
        super(EdgeEnhancement, self).__init__()
        self.conv1 = BasicConv2d(in_channels, out_channels, 3, padding=1)
        self.conv2 = BasicConv2d(out_channels, out_channels, 3, padding=1)
        
        # 使用分组卷积进行边缘检测，每个通道独立处理
        self.edge_conv = nn.Conv2d(out_channels, out_channels, 3, padding=1, 
                                  groups=out_channels, bias=False)
        self.relu = nn.ReLU(inplace=True)
        
        # 初始化边缘检测核
        edge_kernel = torch.tensor([
            [-1, -1, -1],
            [-1,  8, -1],
            [-1, -1, -1]
        ], dtype=torch.float32)
        
        # 为每个通道设置相同的边缘检测核
        weight = edge_kernel.unsqueeze(0).unsqueeze(0).repeat(out_channels, 1, 1, 1)
        self.edge_conv.weight.data = weight
        
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        edge = self.edge_conv(x)
        # 边缘增强：直接返回增强后的特征
        return x + 0.1 * edge

class ContrastEnhancement(nn.Module):
    """对比度增强模块"""
    def __init__(self, in_channels, out_channels):
        super(ContrastEnhancement, self).__init__()
        self.conv1 = BasicConv2d(in_channels, out_channels, 3, padding=1)
        self.conv2 = BasicConv2d(out_channels, out_channels, 3, padding=1)
        self.attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(out_channels, out_channels // 4, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels // 4, out_channels, 1),
            nn.Sigmoid()
        )
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        # 通道注意力增强对比度
        att = self.attention(x)
        return x * att

class UltrasoundRFB(nn.Module):
    """针对超声图像优化的RFB模块"""
    def __init__(self, in_channel, out_channel):
        super(UltrasoundRFB, self).__init__()
        self.relu = nn.ReLU(True)
        
        # 简化的RFB分支 - 确保空间尺寸一致
        self.branch0 = nn.Sequential(
            BasicConv2d(in_channel, out_channel, 1),
        )
        self.branch1 = nn.Sequential(
            BasicConv2d(in_channel, out_channel, 1),
            BasicConv2d(out_channel, out_channel, 3, padding=1),
            BasicConv2d(out_channel, out_channel, 3, padding=3, dilation=3)
        )
        self.branch2 = nn.Sequential(
            BasicConv2d(in_channel, out_channel, 1),
            BasicConv2d(out_channel, out_channel, 3, padding=1),
            BasicConv2d(out_channel, out_channel, 3, padding=5, dilation=5)
        )
        self.branch3 = nn.Sequential(
            BasicConv2d(in_channel, out_channel, 1),
            BasicConv2d(out_channel, out_channel, 3, padding=1),
            BasicConv2d(out_channel, out_channel, 3, padding=7, dilation=7)
        )
        
        # 超声图像专用分支
        self.speckle_branch = SpeckleNoiseReduction(in_channel, out_channel)
        self.edge_branch = EdgeEnhancement(in_channel, out_channel)
        self.contrast_branch = ContrastEnhancement(in_channel, out_channel)
        
        # 特征融合
        self.conv_cat = BasicConv2d(7*out_channel, out_channel, 3, padding=1)
        self.conv_res = BasicConv2d(in_channel, out_channel, 1)
        
        # 自适应权重学习
        self.weight_conv = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(7*out_channel, 7, 1),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        # 原始RFB分支
        x0 = self.branch0(x)
        x1 = self.branch1(x)
        x2 = self.branch2(x)
        x3 = self.branch3(x)
        
        # 超声图像专用分支
        x_speckle = self.speckle_branch(x)
        x_edge = self.edge_branch(x)
        x_contrast = self.contrast_branch(x)
        # 特征拼接
        x_cat = torch.cat((x0, x1, x2, x3, x_speckle, x_edge, x_contrast), 1)
        # 自适应权重融合
        weights = self.weight_conv(x_cat)
        weighted_features = []
        for i in range(7):
            weighted_features.append(x_cat[:, i*self.conv_cat.conv.in_channels//7:(i+1)*self.conv_cat.conv.in_channels//7] * weights[:, i:i+1])
        x_weighted = torch.cat(weighted_features, dim=1)
        
        # 最终融合
        x_fused = self.conv_cat(x_weighted)
        x_res = self.conv_res(x)
        
        return self.relu(x_fused + x_res)

class AdaptiveRFB(nn.Module):
    """自适应RFB模块，根据输入特征动态调整分支权重"""
    def __init__(self, in_channel, out_channel):
        super(AdaptiveRFB, self).__init__()
        self.ultrasound_rfb = UltrasoundRFB(in_channel, out_channel)
        
        # 输入特征分析模块
        self.feature_analyzer = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channel, 64, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 32, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 3, 1),  # 3个特征指标：噪声水平、边缘强度、对比度
            nn.Sigmoid()
        )
        
        # 分支选择门控
        self.branch_gate = nn.Sequential(
            nn.Conv2d(3, 7, 1),  # 7个分支的权重
            nn.Softmax(dim=1)
        )
        
    def forward(self, x):
        # 分析输入特征
        feature_indicators = self.feature_analyzer(x)
        branch_weights = self.branch_gate(feature_indicators)
        
        # 应用RFB
        rfb_output = self.ultrasound_rfb(x)
        
        # 根据特征指标调整输出
        return rfb_output * (1 + feature_indicators.mean(dim=1, keepdim=True)) 